package com.gaofans.nrpc.codegen.processor;

import com.gaofans.nrpc.codegen.annontion.ProxyGen;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.Elements;
import javax.lang.model.util.Types;
import javax.tools.JavaFileObject;
import java.io.IOException;
import java.io.Writer;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;

/**
 * 为标记@ProxyGen的接口生成代理类
 * @author GaoFans
 */
public class CodegenProcessor extends AbstractProcessor {

    private Types typeUtils;
    private Elements elementUtils;
    private Filer filer;
    private Messager messager;
    private final static Logger LOGGER = LoggerFactory.getLogger(CodegenProcessor.class);
    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        elementUtils = processingEnv.getElementUtils();
        typeUtils = processingEnv.getTypeUtils();
        filer = processingEnv.getFiler();
        messager = processingEnv.getMessager();
    }

    @Override
    public SourceVersion getSupportedSourceVersion(){
        return SourceVersion.latestSupported();
    }

    @Override
    public Set<String> getSupportedAnnotationTypes(){
        Set<String> annotations = new LinkedHashSet<>();
        annotations.add(ProxyGen.class.getCanonicalName());
        return annotations;
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        Set<? extends Element> elements = roundEnv.getElementsAnnotatedWith(ProxyGen.class);
        // 遍历所有被注解了@ProxyGen的元素
        for (Element annotatedElement : elements) {
            // 检查被注解为@ProxyGen的元素是否是一个接口
            if (annotatedElement.getKind() != ElementKind.INTERFACE) {
                throw new IllegalStateException("使用@ProxyGen标记的类必须是一个接口");
            }
            //判断是否依赖了client包，有的话生成serviceProxy
            if (classExisted("com.gaofans.nrpc.client.bootstrap.ChannelManager")){
                proxyBuild(annotatedElement);
            }
            //判断是否依赖了server包，有的话生成serviceHandler
            if (classExisted("com.gaofans.nrpc.server.service.ServiceHandler")){
                serviceHandlerBuild(annotatedElement);
            }

        }
        return true;
    }

    private void proxyBuild(Element classElement){
        TypeElement te = (TypeElement) classElement;
        List<? extends TypeParameterElement> typeParameters = te.getTypeParameters();
        if (typeParameters != null && typeParameters.size() > 0){
            throw new IllegalStateException("不支持使用@ProxyGen标记的类使用泛型");
        }
        List<? extends TypeMirror> interfaces = te.getInterfaces();
        if (interfaces != null && interfaces.size() > 0){
            throw new IllegalStateException("不支持使用@ProxyGen标记的类继承接口");
        }
        Element up = te.getEnclosingElement();
        if (up.getKind().equals(ElementKind.PACKAGE)){
            PackageElement pe = (PackageElement) up;
            String packageName = pe.getQualifiedName().toString();
            String className = te.getSimpleName().toString();
            String newClassName = className + "GenRpcProxy";
            StringBuilder builder = new StringBuilder()
                    .append("package " + packageName + ";\n\n")
                    .append("import com.gaofans.nrpc.client.bootstrap.ChannelManager;\n" +
                            "import com.gaofans.nrpc.common.context.ContextRequest;\n" +
                            "import com.gaofans.nrpc.common.id.IdGenerator;\n" +
                            "import com.gaofans.nrpc.common.processor.MultiSink;\n" +
                            "import com.gaofans.nrpc.common.processor.SimpleEventListener;\n" +
                            "import com.gaofans.nrpc.common.processor.SingleSink;\n" +
                            "import com.gaofans.nrpc.common.protocol.RpcRequest;\n" +
                            "import com.gaofans.nrpc.common.protocol.RpcRequestBuilder;\n" +
                            "import io.netty.channel.Channel;\n" +
                            "import reactor.core.publisher.Flux;\n" +
                            "import reactor.core.publisher.Mono;\n" +
                            "import java.net.URI;\n")
                    .append("public class ")
                    .append(newClassName)
                    .append(" implements ")
                    .append(className)
                    .append(" {\n\n")
                    .append("   private final ChannelManager channelManager;\n" +
                            "   private final URI uri;\n" +
                            "   private final IdGenerator idGenerator;\n")
                    .append("   public "+newClassName+"(ChannelManager channelManager, URI uri, IdGenerator idGenerator) {\n" +
                            "       this.channelManager = channelManager;\n" +
                            "       this.uri = uri;\n" +
                            "       this.idGenerator = idGenerator;\n" +
                            "   }\n");
            List<? extends Element> enclosedElements = te.getEnclosedElements();
            for (Element element : enclosedElements) {
                if (element.getKind().equals(ElementKind.METHOD)){
                    ExecutableElement method = (ExecutableElement) element;
                    TypeMirror returnType = method.getReturnType();
                    if (returnType.getKind().equals(TypeKind.DECLARED)){
                        DeclaredType rt = (DeclaredType) returnType;
                        Element rm = rt.asElement();
                        if (!isReactor(rm.toString())){
                            throw new IllegalStateException("方法返回值必须是Flux或Mono");
                        }
                        List<? extends TypeMirror> typeArguments = rt.getTypeArguments();
                        String typeName = null;
                        if (typeArguments != null && typeArguments.size() > 0){
                            TypeMirror typeMirror = typeArguments.get(0);
                            if (typeMirror.getKind().equals(TypeKind.DECLARED)){
                                typeName = ((DeclaredType) (typeMirror)).asElement().toString();
                            }else {
                                throw new IllegalStateException("方法泛型异常");
                            }
                        }
                        builder.append("    public ");
                        builder.append(rt).append(" ");
                        builder.append(method.getSimpleName());
                        builder.append("(");
                        List<? extends VariableElement> parameters = method.getParameters();
                        StringBuilder tp = new StringBuilder("new Class[]{");
                        StringBuilder tv = new StringBuilder("new Object[]{");
                        for (int i = 0; i < parameters.size(); i++) {
                            builder.append(parameters.get(i).asType()).append(" ");
                            builder.append(parameters.get(i).getSimpleName());
                            tp.append(((DeclaredType)parameters.get(i).asType()).asElement()).append(".class");
                            tv.append(parameters.get(i).getSimpleName());
                            if (i != parameters.size()-1){
                                builder.append(",");
                                tp.append(",");
                                tv.append(",");
                            }
                        }
                        tp.append("}");
                        tv.append("}");
                        builder.append(")");
                        builder.append("{\n");
                        builder.append("        Channel channel = channelManager.connect(uri);\n" +
                                "        RpcRequest request = new RpcRequestBuilder()\n" +
                                "                .id(idGenerator.nextId())\n" +
                                "                .className(\""+te.getQualifiedName()+"\")\n" +
                                "                .methodName(\""+method.getSimpleName()+"\")\n" +
                                "                .typeParameters("+tp+")\n" +
                                "                .parametersVal("+tv+")\n" +
                                "                .build();\n" +
                                "        ContextRequest contextRequest = new ContextRequest.ContextRequestBuilder(request).executor(channelManager.chooser()).build();\n" +
                                "        return "+ rm +".create(sink -> {\n" +
                                "            channelManager.processor().register(contextRequest.getSequenceId(), new SimpleEventListener(new "+("reactor.core.publisher.Flux".equals(rm.toString())?"MultiSink":"SingleSink")+"(sink)));\n" +
                                "            try {\n" +
                                "                channel.writeAndFlush(contextRequest);\n" +
                                "            }catch (Exception e){\n" +
                                "                channelManager.processor().unRegister(contextRequest.getSequenceId());\n" +
                                "            }\n" +
                                "        })");
                        builder.append(";\n");
                        builder.append("    }\n");
                    }else {
                        throw new IllegalStateException("方法返回值必须是Flux或Mono");
                    }
                }
            }
            builder.append("}");

            writeFile(packageName, newClassName, builder);
            LOGGER.info("{}类proxy生成成功",te.getQualifiedName());
        }
    }

    private void serviceHandlerBuild(Element classElement){
        TypeElement te = (TypeElement) classElement;
        Element up = te.getEnclosingElement();
        if (up.getKind().equals(ElementKind.PACKAGE)){
            PackageElement pe = (PackageElement) up;
            String packageName = pe.getQualifiedName().toString();
            String className = te.getSimpleName().toString();
            char[] chars = className.toCharArray();
            chars[0] = Character.toLowerCase(chars[0]);;
            String classNameVar = String.valueOf(chars);
            String newClassName = className + "GenServiceHandler";
            StringBuilder builder = new StringBuilder().append("package ").append(packageName).append(";\n\n");
            builder.append("import com.gaofans.nrpc.common.protocol.RpcRequest;\n" +
                    "import com.gaofans.nrpc.common.protocol.RpcResponse;\n" +
                    "import com.gaofans.nrpc.server.service.ServiceHandler;\n" +
                    "import io.netty.channel.Channel;\n" +
                    "import org.slf4j.Logger;\n" +
                    "import org.slf4j.LoggerFactory;\n");
            builder.append("public class ").append(newClassName).append(" implements ServiceHandler {\n\n");
            builder.append("    private final ").append(te.getQualifiedName()).append(" ").append(classNameVar).append(";\n");
            builder.append("    private final static Logger LOGGER = LoggerFactory.getLogger(").append(newClassName).append(".class);\n");
            builder.append("    public "+newClassName+"("+className+" "+classNameVar+") {\n" +
                            "        this."+classNameVar+" = "+classNameVar+";\n" +
                            "    }\n");
            builder.append("    @Override\n" +
                    "    public void handle(RpcRequest request, Channel channel) {\n");
            builder.append("        String methodName = request.getMethodName();\n");
            builder.append("        switch (methodName) {\n");
            List<? extends Element> elements = te.getEnclosedElements();
            for (Element element : elements) {
                ExecutableElement method = (ExecutableElement) element;
                TypeMirror returnType = method.getReturnType();
                if (returnType.getKind().equals(TypeKind.DECLARED)) {
                    DeclaredType rt = (DeclaredType) returnType;
                    Element rm = rt.asElement();
                    if ("reactor.core.publisher.Flux".equals(rm.toString())){
                        fluxCaseBuild(classNameVar, builder, method);
                    }else if ("reactor.core.publisher.Mono".equals(rm.toString())){
                        monoCaseBuild(classNameVar, builder, method);
                    }else {
                        throw new IllegalStateException("方法返回值必须是Flux或Mono");
                    }

                }else {
                    throw new IllegalStateException("方法返回值必须是Flux或Mono");
                }

            }
            builder.append("            default:\n" +
                    "                   throw new IllegalArgumentException(\"method not found\");\n");
            builder.append("        }\n");
            builder.append("    }\n");
            builder.append("}");
            writeFile(packageName, newClassName, builder);
            LOGGER.info("{}类ServiceHandler生成成功",te.getQualifiedName());
        }
    }

    private void monoCaseBuild(String classNameVar, StringBuilder builder, ExecutableElement method) {
        String methodName = method.getSimpleName().toString();
        builder.append("        case \"").append(methodName).append("\":\n");
        StringBuilder requestParam = getRequestParam(method);
        builder.append("            try{\n");
        builder.append("                ").append(classNameVar).append(".").append(methodName).append("(").append(requestParam).append(")\n");
        builder.append("                    .doOnSuccess(res -> {\n");
        builder.append("                        RpcResponse response = new RpcResponse();\n" +
                "                               response.setSequenceId(request.getSequenceId());\n" +
                "                               response.setResult(res);\n" +
                "                               response.setStream(false);\n" +
                "                               response.setSuccess(true);\n" +
                "                               response.setEnd(true);\n" +
                "                               channel.writeAndFlush(response);\n");
        builder.append("                    }).doOnError(e -> {\n");
        builder.append("                        RpcResponse response = new RpcResponse();\n" +
                "                               response.setSequenceId(request.getSequenceId());\n" +
                "                               response.setSuccess(false);\n" +
                "                               response.setStream(false);\n" +
                "                               response.setError(((java.lang.Throwable)e).getMessage());\n" +
                "                               response.setEnd(true);\n" +
                "                               channel.writeAndFlush(response);\n                  }).subscribe();\n");
        builder.append("            }catch (java.lang.Exception e){\n");
        builder.append("                RpcResponse response = new RpcResponse();\n" +
                        "                response.setSequenceId(request.getSequenceId());\n" +
                        "                response.setSuccess(false);\n" +
                        "                response.setStream(false);\n" +
                        "                response.setError(e.getMessage());\n" +
                        "                response.setEnd(true);\n" +
                        "                channel.writeAndFlush(response);\n");
        builder.append("            }\n");
        builder.append("                break;\n");
    }

    private void fluxCaseBuild(String classNameVar, StringBuilder builder, ExecutableElement method) {
        String methodName = method.getSimpleName().toString();
        builder.append("        case \"").append(methodName).append("\":\n");
        StringBuilder requestParam = getRequestParam(method);
        builder.append("            try{\n");
        builder.append("                ").append(classNameVar).append(".").append(methodName).append("(").append(requestParam).append(")\n");
        builder.append("                    .doOnComplete(() -> {\n" +
                "                            RpcResponse response = new RpcResponse();\n" +
                "                            response.setSequenceId(request.getSequenceId());\n" +
                "                            response.setStream(true);\n" +
                "                            response.setSuccess(true);\n" +
                "                            response.setEnd(true);\n" +
                "                            LOGGER.info(\"request {} complete\",request.getSequenceId());\n" +
                "                            channel.writeAndFlush(response);\n" +
                "                        }).doOnError(e -> {\n" +
                "                            RpcResponse response = new RpcResponse();\n" +
                "                            response.setSequenceId(request.getSequenceId());\n" +
                "                            response.setStream(true);\n" +
                "                            response.setSuccess(false);\n" +
                "                            response.setError(((java.lang.Throwable)e).getMessage());\n" +
                "                            response.setEnd(true);\n" +
                "                            channel.writeAndFlush(response);\n" +
                "                        }).subscribe(res -> {\n" +
                "                            RpcResponse response = new RpcResponse();\n" +
                "                            response.setSequenceId(request.getSequenceId());\n" +
                "                            response.setStream(true);\n" +
                "                            response.setSuccess(true);\n" +
                "                            response.setEnd(false);\n" +
                "                            response.setResult(res);\n" +
                "                            LOGGER.info(\"response {} send {}\",request.getSequenceId(),res);\n" +
                "                            channel.writeAndFlush(response);\n" +
                "                        });\n");
        builder.append("            }catch (java.lang.Exception e){\n");
        builder.append("                RpcResponse response = new RpcResponse();\n" +
                "                response.setSequenceId(request.getSequenceId());\n" +
                "                response.setStream(true);\n" +
                "                response.setSuccess(false);\n" +
                "                response.setError(e.getMessage());\n" +
                "                response.setEnd(true);\n" +
                "                channel.writeAndFlush(response);\n");
        builder.append("            }\n");
        builder.append("                break;\n");
    }

    private StringBuilder getRequestParam(ExecutableElement method) {
        List<? extends VariableElement> parameters = method.getParameters();
        StringBuilder requestParam = new StringBuilder("");
        for (int i = 0; i < parameters.size(); i++) {
            requestParam
                    .append("(")
                    .append(((DeclaredType)parameters.get(i).asType()).asElement()).append(")")
                    .append("request.getParametersVal()[")
                    .append(i)
                    .append("]");
            if (i != parameters.size() - 1){
                requestParam.append(",");
            }
        }
        return requestParam;
    }

    private void writeFile(String packageName, String newClassName, StringBuilder builder) {
        try {
            JavaFileObject source = filer.createSourceFile(packageName +"."+ newClassName);
            Writer writer = source.openWriter();
            writer.write(builder.toString());
            writer.flush();
            writer.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private static boolean classExisted(String className){
        try {
            Class.forName(className);
            return true;
        } catch (ClassNotFoundException e) {
            return false;
        }
    }

    private static boolean isReactor(String name){
        return "reactor.core.publisher.Flux".equals(name) || "reactor.core.publisher.Mono".equals(name);
    }
}
